from logging import exception
from socket import timeout
import threading
from USocket import UnreliableSocket, sockets
from socket import timeout
from threading import Thread
import time
import queue
from prettytable import PrettyTable

# import copy

LISTEN = 0
SYN_RECV = 1
ESTABLISH = 2
TIMEOUT = 5
ALPHA = 0.15
BETA = 0.25
WIN_SIZE = 10


class RDTSocket(UnreliableSocket):
    """
    The functions with which you are to build your RDT.
    -   recvfrom(bufsize)->bytes, addr
    -   sendto(bytes, address)
    -   bind(address)

    You can set the mode of the socket.
    -   settimeout(timeout)
    -   setblocking(flag)
    By default, a socket is created in the blocking mode.
    https://docs.python.org/3/library/socket.html#socket-timeouts

    """

    def __init__(self, rate=None, debug=True):
        super().__init__(rate=rate)
        self._rate = rate
        self._send_to = None
        self._recv_from = None
        self.debug = debug
        #############################################################################
        # TODO: ADD YOUR NECESSARY ATTRIBUTES HERE
        #############################################################################
        self.port = 49152
        self.window_size = 2
        self.send_buffer = queue.Queue()  # send() put in
        self.recv_buffer = bytes()  # receive() take out
        self.recv_ack_buffer = queue.Queue()  # receive()
        self._seq = 0
        self._ack = 0
        self.closed = False
        self.estimatedRTT = 1
        self.devRTT = 1
        self.timeout = 10
        self.name = ''
        self.lastReceiveTime = None
        self.lock = None
        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################

    def accept(self) -> ('RDTSocket', (str, int)):
        """
        Accept a connection. The socket must be bound to an address and listening for
        connections. The return value is a pair (conn, address) where conn is a new
        socket object usable to send and receive data on the connection, and address
        is the address bound to the socket on the other end of the connection.

        This function should be blocking.

        先bind一个地址，地址是指本机有
        多个ip地址时的情况，留空即所有地址
        然后监听这个接口，准备接收链接
        当有信息发送过来进行处理以及建立链接
        """
        conn, addr = RDTSocket(self._rate), None
        #############################################################################
        # TODO: YOUR CODE HERE                                                      #
        #############################################################################
        # listen是tcp对象特有的功能，而我们这里建立的相当于是udp对象
        # sockets[id(conn)].listen(0)#为啥这句不行呢？？！！
        self.name = 'server'
        conn.name = 'conn'
        # TODO: bind an address first
        conn.bind(('', 10001))
        # init seq_num
        conn._seq = 0
        conn._ack = 0
        # state = LISTEN
        while True:
            # recvfrom running in blocking mode
            data, addr = self.recvfrom(4096)
            rdtSegment = RDTSegment.decode(data)
            RDTSegment.print(rdtSegment)
            with open(file=self.name+'_'+'receive_log.txt', mode='w') as f:
                RDTSegment.print(rdtSegment, file=f)

            if rdtSegment.syn:
                # state = SYN_RECV
                conn._ack = rdtSegment.seq_num
                reply_segment = RDTSegment(
                    seq_num=conn._seq, ack_num=conn._ack, syn=1, fin=0, ack=1)
                conn.sendto(reply_segment.encode(), addr)
                RDTSegment.print(rdtSegment)
                with open(file=conn.name+'_'+'send_log.txt', mode='w') as f:
                    RDTSegment.print(rdtSegment, file=f)
                conn.settimeout(TIMEOUT)
            else:
                continue

            # recvfrom running in blocking mode with timeout set
            # if time out will return to listen state
            try:
                data, addr = conn.recvfrom(4096)
                rdtSegment = RDTSegment.decode(data)
                RDTSegment.print(rdtSegment)
                with open(file=conn.name+'_'+'receive_log.txt', mode='w') as f:
                    RDTSegment.print(rdtSegment, file=f)
                if rdtSegment.ack and rdtSegment.ack_num == conn._seq:
                    # state = ESTABLISH
                    conn._send_to = addr
                    conn._recv_from = addr
                    conn._seq += 1
                    conn._ack = rdtSegment.seq_num
                    break
            except timeout:
                # state = LISTEN
                # init seq_num
                conn._seq = 0
                conn._ack = 0
                conn.settimeout(None)
            except:
                conn._seq = 0
                conn._ack = 0
                conn.settimeout(None)
        # conn.setblocking(0)
        # conn.settimeout(None)
        conn.lock = threading.Lock()
        conn.send_thread = Thread(target=thread_send, args=(conn, conn.lock))
        conn.send_thread.setDaemon(True)
        conn.recv_thread = Thread(target=thread_recv, args=(conn, conn.lock))
        conn.recv_thread.setDaemon(True)
        conn.send_thread.start()
        conn.recv_thread.start()
        time.sleep(1)
        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################
        return conn, addr

    def connect(self, address: (str, int)):
        """
        Connect to a remote socket at address.
        Corresponds to the process of establishing a connection on the client side.

        首先通过network向远端发送信息
        如果远端回复信息，就建立连接
        如果远端未回复说明建立失败

        用三次握手？
        定义一个包segment，里面加上flag
        SYN seq=0
        SYN ack=1, ACK seq=0
        ACK seq=0, ack=1
        """
        #############################################################################
        # TODO: YOUR CODE HERE                                                      #
        #############################################################################
        self.name = 'client'
        # TODO: bind an address first
        self.bind(('', 10000))
        # sockets[id(self)] = sockets.get()

        # init seq_num
        self._seq = 0
        self._ack = 0
        # retry 3 times
        i = 0
        while i < 3:
            rdtSegment = RDTSegment(seq_num=self._seq, ack_num=self._ack, length=0,
                                    syn=1, fin=0, ack=0)
            self.sendto(rdtSegment.encode(), address)

            RDTSegment.print(rdtSegment)
            with open(file=self.name+'_'+'send_log.txt', mode='w') as f:
                RDTSegment.print(rdtSegment, file=f)

            # print('trying to send to ', address)
            try:
                self.settimeout(TIMEOUT)
                data, addr = self.recvfrom(4096)
                # print('get ', data, ' from ', addr)
                rdtSegment = RDTSegment.decode(data)

                RDTSegment.print(rdtSegment)
                with open(file=self.name+'_'+'receive_log.txt', mode='w') as f:
                    RDTSegment.print(rdtSegment, file=f)

                if rdtSegment.ack and rdtSegment.syn and rdtSegment.ack_num == self._seq:
                    self.set_send_to(addr)
                    self.set_recv_from(addr)
                    self._seq += 1
                    self._ack = rdtSegment.seq_num
                    rdtSegment = RDTSegment(seq_num=self._seq,
                                            ack_num=self._ack, ack=1)
                    self.sendto(rdtSegment.encode(), addr)

                    RDTSegment.print(rdtSegment)
                    with open(file=self.name+'_'+'send_log.txt', mode='a+') as f:
                        RDTSegment.print(rdtSegment, file=f)

                    self._seq += 1

                    # 在客户端也需设定收发线程吗？
                    # self.setblocking(0)
                    # self.settimeout(None)
                    self.lock = threading.Lock()
                    self.send_thread = Thread(
                        target=thread_send, args=(self, self.lock))
                    self.send_thread.setDaemon(True)
                    self.send_thread.start()
                    self.recv_thread = Thread(
                        target=thread_recv, args=(self, self.lock))
                    self.recv_thread.setDaemon(True)
                    self.recv_thread.start()

                    return
            except:
                i += 1
                continue

        print("connection failed, please retry")

        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################

    def recv(self, bufsize: int) -> bytes:
        """
        Receive data from the socket.
        The return value is a bytes object representing the data received.
        The maximum amount of data to be received at once is specified by bufsize.

        Note that ONLY data send by the peer should be accepted.
        In other words, if someone else sends data to you from another address,
        it MUST NOT affect the data returned by this function.
        只需要将bufsize大小的data返回即可
        """
        data = None
        assert self._recv_from, "Connection not established yet. Use recvfrom instead."
        #############################################################################
        # TODO: YOUR CODE HERE                                                      #
        #############################################################################
        # data = bytearray()
        # expect = 0
        # ack = RDTSegment(seq_num=0, ack_num=expect,
        #                  length=0, payload=None, ack=1, )
        # self.settimeout(socket.TIMEOUT)
        # timeout_count = -1
        # 一直while好像是阻塞模式？
        """
        改到thread里拆分segment，发送ack包以及判断fin
        """
        # while not self.recv_buffer.empty():
        #     rcv_segment = self.recv_buffer.get()
        #     seq = rcv_segment.get_seq_num()
        #     fin = rcv_segment.get_fin()
        #     if not fin:
        #         ack_segment = RDTSegment(seq_num=0, ack_num=seq)
        #         self.send_buffer.put(ack_segment)
        #
        #         data = rcv_segment.get_payload()
        #
        #     else:
        #         pass
        self.lastReceiveTime = time.time()
        while True:
            # self.lock.acquire()
            if self.recv_buffer != b'':
                if len(self.recv_buffer) > bufsize:
                    data = self.recv_buffer[:bufsize]
                    self.recv_buffer = self.recv_buffer[bufsize:]
                else:
                    data = self.recv_buffer
                    self.recv_buffer = bytes()
                # self.lock.release()
                break
            if time.time() - self.lastReceiveTime > 4*self.timeout:
                data = b''
                break
            # self.lock.release()
            # raise Exception('Time out')
        # try:
        #     rcv_data, addr = self.recvfrom(bufsize=bufsize)
        # except timeout:
        #     timeout_count += 1
        #     if timeout_count > 3:
        #         raise timeout
        # else:
        #     if addr == self._recv_from:
        #         timeout_count = 0
        #         break
        #     try:
        #         rdtSegment = RDTSegment.decode(rcv_data)
        #     except:
        #         continue
        #     if rdtSegment.seq_num == expect:
        #         if not rdtSegment.fin:
        #             data.extend(rdtSegment.payload)
        #         else:
        #             # 如果传来fin的信号
        #             # 过程
        #             # c - s
        #             # c: i have sent all - fin
        #             # s: (when receive packet seq num is the expected num) i have received all - fin+ack
        #             # c: ok, connection fin - ack
        #             # connection between c and s closed
        #
        #             break
        #         ack.ack_num = expect
        #         expect = (expect + 1) % self.window_size
        #     self.sendo(ack.encode(), self._recv_from)
        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################
        return data

    def send(self, bytes: bytes):
        """
        Send data to the socket.
        The socket must be connected to a remote socket, i.e. self._send_to must not be none.

        放到一个buffer池里面，buffer池由另一个线程进行发送
        可以在这个地方进行打包
        """
        assert self._send_to, "Connection not established yet. Use sendto instead."
        #############################################################################
        # TODO: YOUR CODE HERE                                                      #
        #############################################################################
        # 每次调用send都需将bytes放入buffer队列？
        # 检查序号则需在recv中进行？
        bytes_len = 1024
        while len(bytes) > 0:
            if len(bytes) > bytes_len:
                payload = bytes[:1024]
                bytes = bytes[1024:]
                segment = RDTSegment(
                    seq_num=self._seq, length=bytes_len, payload=payload)
            else:
                segment = RDTSegment(
                    seq_num=self._seq, length=len(bytes), payload=bytes)
                bytes = b''
            self._seq += 1
            self.send_buffer.put(segment)
            # print('put ',self._seq,' into send_buffer')

        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################

    def close(self):
        """
        Finish the connection and release resources. For simplicity, assume that
        after a socket is closed, neither further sends nor receives are allowed.

        可能要进行关闭的握手
        给的代码的逻辑是没有数据就会关
        """
        #############################################################################
        # TODO: YOUR CODE HERE                                                      #
        #############################################################################
        # TODO: 进行关闭的四次握手
        self._recv_from = None
        self._send_to = None
        self.closed = True
        #############################################################################
        #                             END OF YOUR CODE                              #
        #############################################################################
        super().close()

    def set_send_to(self, send_to):
        self._send_to = send_to

    def set_recv_from(self, recv_from):
        self._recv_from = recv_from

    def calculate_timeout(self, sampleRTT):
        self.estimatedRTT = (1 - ALPHA) * self.estimatedRTT + ALPHA * sampleRTT
        self.devRTT = (1-BETA) * self.devRTT + BETA * \
            abs(self.estimatedRTT-sampleRTT)
        self.timeout = self.estimatedRTT + 4 * self.devRTT


"""
You can define additional functions and classes to do thing such as packing/unpacking packets, or threading.

"""


def thread_send(rdt: RDTSocket, lock: threading.Lock):
    # 我的理解：send thred 一直将send buffer里的数据发给对方，而send方法是把数据放入buffer？
    # 所以这里不应该调用rdt.send（）比较好？
    # TODO: improve by window, selective repeat or go-back N
    while not rdt.closed:
        if not rdt.send_buffer.empty():
            rdtSegment = rdt.send_buffer.get()
            # print('send')
            # RDTSegment.print(rdtSegment)
            expect = rdtSegment.seq_num
            data = rdtSegment.encode()
            t0 = time.time()
            rdt.sendto(data, rdt._send_to)
            with open(file=rdt.name+'_'+'send_log.txt', mode='a+') as f:
                RDTSegment.print(rdtSegment, file=f)
            while True:
                try:
                    # time out can be set to RTT
                    # print(rdt.timeout)
                    recv = rdt.recv_ack_buffer.get(
                        block=True, timeout=rdt.timeout)
                except:
                    # t0 = time.time()
                    rdt.sendto(data, rdt._send_to)
                    with open(file=rdt.name+'_'+'resend_log.txt', mode='a+') as f:
                        RDTSegment.print(rdtSegment, file=f)
                else:
                    if recv.ack and recv.ack_num == expect:
                        rdt.calculate_timeout(time.time()-t0)
                        break
                    else:
                        # t0 = time.time()
                        rdt.sendto(data, rdt._send_to)
                        with open(file=rdt.name+'_'+'resend_log.txt', mode='a+') as f:
                            RDTSegment.print(rdtSegment, file=f)


def thread_recv(rdt: RDTSocket, lock: threading.Lock):
    rdt.setblocking(True)
    t0 = None
    while not rdt.closed:
        # Problem: when to handle the data not from the assume remote socket
        # TODO: need to determine the buffer size here
        try:
            if rdt.closed:
                return
            rdt.settimeout(rdt.timeout)
            data, addr = rdt.recvfrom(4096)
        except OSError:
            return
        except:
            ack = RDTSegment(ack_num=rdt._ack, ack=1)
            # t0 = time.time()
            rdt.sendto(ack.encode(), addr=rdt._send_to)
            with open(file=rdt.name+'_'+'reAck_log.txt', mode='a+') as f:
                print('send', file=f)
                RDTSegment.print(ack, file=f)
        else:
            # if handle here
            rdt.lastReceiveTime = time.time()
            if addr == rdt._recv_from:
                try:
                    rdtSegment = RDTSegment.decode(data)
                    print('receive')
                    RDTSegment.print(rdtSegment)
                    print()
                    with open(file=rdt.name+'_'+'receive_log.txt', mode='a+') as f:
                        print('receive', file=f)
                        RDTSegment.print(rdtSegment, file=f)
                except:
                    ack = RDTSegment(ack_num=rdt._ack, ack=1)
                    # t0 = time.time()
                    rdt.sendto(ack.encode(), addr=rdt._send_to)
                    with open(file=rdt.name+'_'+'corruptAck_log.txt', mode='a+') as f:
                        print('send', file=f)
                        RDTSegment.print(ack, file=f)
                    continue
                else:
                    lock.acquire()
                    # The recv packet is the expected packet
                    if rdtSegment.ack:
                        rdt.recv_ack_buffer.put(rdtSegment)
                    elif rdt._ack + 1 == rdtSegment.seq_num:

                        if t0 != None:
                            rdt.calculate_timeout(time.time()-t0)
                        ack = RDTSegment(ack_num=rdtSegment.seq_num, ack=1)
                        t0 = time.time()
                        rdt.sendto(ack.encode(), addr=rdt._send_to)
                        with open(file=rdt.name+'_'+'ack_log.txt', mode='a+') as f:
                            print('send', file=f)
                            RDTSegment.print(ack, file=f)
                        rdt._ack = rdtSegment.seq_num
                        rdt.recv_buffer += rdtSegment.payload
                    else:
                        ack = RDTSegment(ack_num=rdt._ack, ack=1)
                        # t0 = time.time()
                        rdt.sendto(ack.encode(), addr=rdt._send_to)
                        with open(file=rdt.name+'_'+'corruptAck_log.txt', mode='a+') as f:
                            print('send', file=f)
                            RDTSegment.print(ack, file=f)
                    lock.release()
            # rdt.recv_buffer.put(data)
        # if not, just put the data into the buffer
        # Then we need to handle the unexpected data
        # 1. Throw the data
        # 2. Put the data at the end
        # we can choose 1 to avoid resources waste


class RDTSegment:
    """
    Segment Format:
        syn:int -- 1 byte
        fin:int -- 1 byte
        ack:int -- 1 byte
        seq_num:int(0~2^32-1) -- 4 bytes
        ack_num:int(0~4294967295) -- 4 bytes
        length:int(0~4294967295) -- 4 bytes
        checksum:int(0~65535) -- 2 bytes
        payload:bytes -- length bytes

    """
    HEADER_LEN = 17

    def __init__(self, seq_num=0, ack_num=0, length=0, payload=None,
                 syn: bool = 0, fin: bool = 0, ack: bool = 0):
        self.syn = syn
        self.fin = fin
        self.ack = ack
        self.seq_num = seq_num
        self.ack_num = ack_num
        self.length = length
        # calculate checksum after all data pass
        # self.checksum = checksum
        self.payload = payload

    def get_syn(self):
        return self.syn

    def get_fin(self):
        return self.fin

    def get_ack(self):
        return self.ack

    def get_seq_num(self):
        return self.seq_num

    def get_ack_num(self):
        return self.ack_num

    def get_length(self):
        return self.length

    # def get_checksum(self):
    #     return self.checksum

    def get_data(self):
        return self.payload

    def get_payload(self):
        return self.payload

    def encode(self) -> bytes:
        flag = (self.syn, self.fin, self.ack)
        byte_arr = bytes(flag)
        seq_num_byte = self.seq_num.to_bytes(length=4, byteorder='big')
        ack_num_byte = self.ack_num.to_bytes(length=4, byteorder='big')
        length_byte = self.length.to_bytes(length=4, byteorder='big')
        checksum_byte = RDTSegment.calculate_checksum(
            segment=self).to_bytes(length=2, byteorder='big')
        byte_arr += seq_num_byte
        byte_arr += ack_num_byte
        byte_arr += length_byte
        byte_arr += checksum_byte
        if self.payload:
            byte_arr += self.payload

        # self.checksum = RDTSegment.calculate_checksum(byte_arr)
        # checksum_byte = self.checksum

        return byte_arr

    @staticmethod
    def decode(byte_arr: bytes) -> ('RDTSegment'):
        syn = byte_arr[0]
        fin = byte_arr[1]
        ack = byte_arr[2]
        seq_num = int.from_bytes(byte_arr[3:7], byteorder='big')
        ack_num = int.from_bytes(byte_arr[7:11], byteorder='big')
        length = int.from_bytes(byte_arr[11:15], byteorder='big')
        checksum = int.from_bytes(byte_arr[15:17], byteorder='big')
        payload = byte_arr[17:]

        segment = RDTSegment(syn=syn, fin=fin, ack=ack, seq_num=seq_num,
                             ack_num=ack_num, length=length, payload=payload)
        # 建议这里加上checksum校验
        # 校验失败raise一个exception，让处理部分可以接收
        if checksum != RDTSegment.calculate_checksum(segment=segment):
            raise Exception('checksumError')
        return segment

    @staticmethod
    def calculate_checksum(segment: 'RDTSegment') -> int:
        """
        The checksum field is the 16 bit one’s complement of the one’s complement sum of
        all 16-bit words in the header and text.
        """
        sum = 0
        flag = (segment.syn, segment.fin, segment.ack)
        byte_arr = bytes(flag)
        seq_num_byte = segment.seq_num.to_bytes(length=4, byteorder='big')
        ack_num_byte = segment.ack_num.to_bytes(length=4, byteorder='big')
        length_byte = segment.length.to_bytes(length=4, byteorder='big')
        byte_arr += seq_num_byte
        byte_arr += ack_num_byte
        byte_arr += length_byte
        if segment.payload:
            byte_arr += segment.payload

        if len(byte_arr) % 2 == 1:
            byte_arr += bytes(0)

        for i in range(int(len(byte_arr) / 2)):
            b = byte_arr[2 * i:2 * i + 2]
            sum += int.from_bytes(b, byteorder='big', signed=False)

        # add other handle with checksum
        checksum = 0xffff - ((0xffff & sum) + (sum >> 16))
        checksum = (0xffff & checksum) + (checksum >> 16)
        # print(sum, checksum)

        return checksum

    @staticmethod
    def print(segment: 'RDTSegment', file=None):
        header = 'seq_num ack_num syn ack fin length'.split(' ')
        row = [segment.seq_num, segment.ack_num, segment.syn,
               segment.ack, segment.fin, segment.length]
        pt = PrettyTable()
        pt._set_field_names(header)
        pt.add_row(row=row)

        if file != None:
            print(time.time(), file=file)
            print(pt, file=file)
            print(file=file)
        else:
            print(pt)


if __name__ == "__main__":
    # c = RDTSocket(1)
    # c.accept()

    num = 2333
    rdt = RDTSegment(ack=1, seq_num=12, ack_num=1212,
                     length=2, payload=num.to_bytes(2, 'big'))
    rdt_byte = rdt.encode()
    print('len=', len(rdt_byte))
    rdt_byte_slice = rdt_byte[17:19]
    print('ack=', rdt_byte[2])
    print('payload=', int.from_bytes(rdt_byte_slice, byteorder='big'))
    print('checksum_byte=', int.from_bytes(rdt_byte[15:17], byteorder='big'))
    print('checksum_cal=', RDTSegment.calculate_checksum(rdt))

    back = RDTSegment.decode(rdt_byte)
    print('seq=', back.seq_num)
    print('checksum_back=', RDTSegment.calculate_checksum(back))
